package edu.gatech.cc.liam.core;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Random;

public class Tensor {
	private double[] values;
	private int[] dimensions;
	
	private int[] dimMults;

	public Tensor(int ... dims) {
		initTensor(dims);
	}
	
	public Tensor(Object[] t) {
		ArrayList<Integer> dimensions = new ArrayList<Integer>();
		dimensions.add(t.length);
		Object[] temp = t;
		while(temp.length > 0 && temp[0].getClass().isArray()) {
			temp = (Object[]) temp[0];
			dimensions.add(temp.length);
		}
		int[] dims = new int[dimensions.size()];
		for(int i=0; i<dims.length; i++)
			dims[i] = dimensions.get(i);
		initTensor(dims);
		fill(t);
	}
	
	public void fill(Object[] t) {
		fillHelper(t, new int[0]);
	}
	private void fillHelper(Object[] t, int[] dims) {
		assert t.length > 0;
		for(int i=0; i<t.length; i++) {
			int[] newD = appendToIntArray(dims, i);
			if(t[i].getClass().isArray()) {
				assert(t[i] instanceof Object[]);
				fillHelper((Object[])t[i], newD);
			} else {
				set(((Number)t[i]).doubleValue(), newD);
			}
		}
	}
	private int[] appendToIntArray(int[] a, int n) {
		int[] r = Arrays.copyOf(a, a.length+1);
		r[a.length] = n;
		return r;
	}
	
	
	private void initTensor(int[] dims) {
		dimensions = dims;
		dimMults = new int[dims.length];
		dimMults[dimMults.length-1] = 1;
		for(int i=1; i<dimMults.length; i++) {
			dimMults[dimMults.length-i-1] = dimMults[dimMults.length-i] * dims[dimMults.length-i];
		}
		int totalSize = dimMults[0]*dims[0];
		values = new double[totalSize];
	}
	
	public int[] getDimensions() {
		return dimensions;
	}

	private int getValueLoc(int ... pos){
		assert pos.length == dimensions.length;
		int loc = 0;
		for(int i=0; i<dimMults.length; i++) {
			loc += pos[i] * dimMults[i];
		}
		return loc;
	}
	
	public double get(int ... pos) {
		return values[getValueLoc(pos)];
	}
	public void set(double value, int ... pos) {
		values[getValueLoc(pos)] = value;
	}
	public Double[] getBottomArray(int ... pos) {
		Double[] rVals = new Double[dimensions[dimensions.length-1]];
		for(int i=0; i<rVals.length; i++) {
			rVals[i] = get(appendToIntArray(pos, i));
		}
		return rVals;
	}

	public void fillWithRandom(Random rand) {
		for(int i=0; i<values.length; i++) {
			values[i] = rand.nextDouble();
		}
	}
	
	@Override
	public String toString() {
		return makeStr(new int[0]);
	}
	
	private String makeStr(int[] startingPos) {
		if(startingPos.length == dimensions.length)
			return Double.toString(get(startingPos));
		int[] newPos = Arrays.copyOf(startingPos, startingPos.length + 1);
		String rStr = "{";
		for(int i=0; i<dimensions[startingPos.length]; i++) {
			newPos[startingPos.length] = i;
			rStr += makeStr(newPos);
			if(i != dimensions[startingPos.length]-1)
				rStr += ", ";
		}
		rStr += "}";
		return rStr;
	}
	
	public static void main(String[] args) {
		Integer[][][] t1 = {{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}};
		Double[][][] t2 = {{{1.0, 2.0, 2.5}, {3.0, 4.0, 4.5}}, {{5.0, 6.0, 6.5}, {7.0, 8.0, 8.5}}};
		
		Tensor T1 = new Tensor(2,2,2);
		System.out.println(Arrays.toString(T1.dimMults));
		for(int i=0; i<t1.length; i++)
			for(int j=0; j<t1[i].length; j++)
				for(int k=0; k<t1[i][j].length; k++) {
					T1.set(t1[i][j][k], i,j,k);
				}
		System.out.println(Arrays.deepToString(t1));
		System.out.println(Arrays.toString(T1.values));
		System.out.println(T1);
		
		Tensor T2 = new Tensor(2,2,3);
		System.out.println(Arrays.toString(T2.dimMults));
				
		for(int i=0; i<t2.length; i++)
			for(int j=0; j<t2[i].length; j++)
				for(int k=0; k<t2[i][j].length; k++) {
					T2.set(t2[i][j][k], i,j,k);
				}
		System.out.println(Arrays.deepToString(t2));
		System.out.println(Arrays.toString(T2.values));
		System.out.println(T2);
		
		
	}
}
